# ==== forwarding_unit.py ====
class ForwardingUnit:
    @staticmethod
    def forward_rs(id_ex, ex_mem, mem_wb):
        if id_ex.rs == 0:
            return id_ex.read_data1, 0

        # EX阶段前递
        if ex_mem.valid and ex_mem.control_signals[
            "RegWrite"] and ex_mem.write_reg != 0 and ex_mem.write_reg == id_ex.rs:
            return ex_mem.ALU_result, 1

        # MEM阶段前递
        if mem_wb.valid and mem_wb.control_signals[
            "RegWrite"] and mem_wb.write_reg != 0 and mem_wb.write_reg == id_ex.rs:
            return mem_wb.ALU_result, 2

        return id_ex.read_data1, 0

    @staticmethod
    def forward_rt(id_ex, ex_mem, mem_wb):
        if id_ex.rt == 0:
            return id_ex.read_data2, 0

        # EX阶段前递
        if ex_mem.valid and ex_mem.control_signals[
            "RegWrite"] and ex_mem.write_reg != 0 and ex_mem.write_reg == id_ex.rt:
            return ex_mem.ALU_result, 1

        # MEM阶段前递
        if mem_wb.valid and mem_wb.control_signals[
            "RegWrite"] and mem_wb.write_reg != 0 and mem_wb.write_reg == id_ex.rt:
            return mem_wb.ALU_result, 2

        return id_ex.read_data2, 0